import os
import cv2
import numpy as np
from enum import Enum, IntEnum, auto
from tqdm import tqdm

class Data(Enum):
    image=0
    depth=1
    mask=2
    vis=3
    cplist=4
    cpbbox=5

dir_list=["image","mask","vis","cplist","cpbbox","cpimask"]

"""
切分pcbsyn3.py生成的数据和bbox
"""
def get_cropped_bboxes(original_bboxes, subimage_coords):
    """
    获取切分后小图片中的边界框

    :param original_bboxes: 原图中的边界框列表，每个边界框为 [x1, y1, x2, y2]
    :param subimage_coords: 小图片四个角在原图中的坐标 [x1, y1, x2, y2]
    :return: 切分后小图片中的边界框列表
    """
    subimage_x1, subimage_y1, subimage_x2, subimage_y2 = subimage_coords
    cropped_bboxes = []

    for bbox in original_bboxes:
        c,bbox_x1, bbox_y1, bbox_x2, bbox_y2 = bbox

        # 检查边界框是否与小图片有交集
        if (bbox_x1 < subimage_x2 and bbox_x2 > subimage_x1 and
                bbox_y1 < subimage_y2 and bbox_y2 > subimage_y1):
            # 计算裁剪后的边界框坐标
            cropped_x1 = max(bbox_x1, subimage_x1) - subimage_x1
            cropped_y1 = max(bbox_y1, subimage_y1) - subimage_y1
            cropped_x2 = min(bbox_x2, subimage_x2) - subimage_x1
            cropped_y2 = min(bbox_y2, subimage_y2) - subimage_y1

            # 确保裁剪后的边界框有效
            if cropped_x2 > cropped_x1 and cropped_y2 > cropped_y1:
                cropped_bboxes.append([c,cropped_x1, cropped_y1, cropped_x2, cropped_y2])

    return cropped_bboxes

def get_cropped_bboxes_xywh(original_bboxes,subimage_coords,img_wh,slice_wh):
    iw,ih=img_wh
    sw,sh=slice_wh
    cpbbox_xyxy_list=[(c,(x-w/2)*iw,(y-h/2)*ih,(x+w/2)*iw,(y+h/2)*ih) for c,x,y,w,h in original_bboxes]
    cropped_bboxes_xyxy=get_cropped_bboxes(cpbbox_xyxy_list,subimage_coords)
    cropped_bboxes_xywh=[(int(c),((x1+x2)/2)/sw,((y1+y2)/2)/sh,(x2-x1)/sw,(y2-y1)/sh) for c,x1,y1,x2,y2 in cropped_bboxes_xyxy]
    return cropped_bboxes_xywh

def read_yolo_bbox_txt(file_path):
    bbox_list = []
    try:
        with open(file_path, 'r') as file:
            for line in file:
                # 去除行尾的换行符，并按空格分割字符串
                data = line.strip().split()
                # 将字符串转换为浮点数
                data = [float(x) for x in data]
                bbox_list.append(data)
    except FileNotFoundError:
        print(f"错误：文件 {file_path} 未找到。")
    except Exception as e:
        print(f"发生未知错误：{e}")
    return bbox_list

def split_images_and_masks(data_root, output_folder, x, y, overlap):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    for dir_name in dir_list:
        if not os.path.exists(os.path.join(data_root,dir_name)):
            print("data folder not exist",dir_name)
            return
        if not os.path.exists(os.path.join(output_folder,dir_name)):
                os.makedirs(os.path.join(output_folder,dir_name))

    cnt=0
    big_image_folder_path = os.path.join(data_root,"image")
    output_image_folder=os.path.join(output_folder,'image')
    output_mask_folder=os.path.join(output_folder,'mask')
    output_yololabel_folder=os.path.join(output_folder,'cpbbox')
    output_cpimask_folder=os.path.join(output_folder,'cpimask')
    for root, dirs, files in os.walk(big_image_folder_path):
        for file in tqdm(files):
            if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
                # cnt+=1
                # if cnt>=4:
                #     return 0
                image_path=os.path.join(big_image_folder_path,file)
                mask_path=image_path.replace("image","mask").replace('.png','_mask.png')
                cpimask_path=image_path.replace("image","cpimask").replace('.png','_cpimask.png')
                image = cv2.imread(image_path,cv2.IMREAD_UNCHANGED)
                mask = cv2.imread(mask_path)
                cpimask=cv2.imread(cpimask_path)
                cpbbox_list=read_yolo_bbox_txt(image_path.replace("image","cpbbox").replace('.png','_cpbbox.txt'))
                
                if image is None or mask is None:
                    print(f"Failed to read {image_path} or {mask_path}")
                    continue

                height, width = image.shape[:2]
                start_x = 0
                while start_x + x <= width:
                    start_y = 0
                    while start_y + y <= height:
                        sub_image = image[start_y:start_y + y, start_x:start_x + x]
                        sub_mask = mask[start_y:start_y + y, start_x:start_x + x]
                        sub_cpimask=cpimask[start_y:start_y + y, start_x:start_x + x]
                        sub_yolo_labels=get_cropped_bboxes_xywh(cpbbox_list,(start_x,start_y,
                                                                             (start_x + x),(start_y + y)),(height, width),(x,y))
                        
                        image_name = f"{os.path.splitext(file)[0]}_{start_x}_{start_y}_{x}_{y}_{overlap}.png"
                        image_output_path = os.path.join(output_image_folder, image_name)
                        mask_output_path = os.path.join(output_mask_folder, image_name)
                        cpimask_output_path = os.path.join(output_cpimask_folder, image_name)
                        yololabel_output_path = os.path.join(output_yololabel_folder, image_name.replace('png','txt'))
                        
                        cv2.imwrite(image_output_path, sub_image)
                        cv2.imwrite(mask_output_path, sub_mask)
                        cv2.imwrite(cpimask_output_path, sub_cpimask)

                        with open(yololabel_output_path, 'w') as f2:
                            for label in sub_yolo_labels:
                                bbox_label=label[:5]
                                label_str = " ".join(str(i) for i in bbox_label)
                                f2.write(label_str + "\n")

                        start_y += y - overlap
                    start_x += x - overlap

if __name__ == "__main__":
    # data_root="synbboxpcb/"
    data_root="/home/lhx/pcbsyn/data_zoo/truecolorseg"
    # output_folder="data_zoo/splitsynbboxpcb/"
    output_folder="data_zoo/splitsyntruecolorseg"

    data_root="/home/lhx/pcbsyn/synpcbwithjrsdepbg"
    output_folder="data_zoo/splitsynpcbwithjrsdepbg"
    x = 1024  # 小块边长x
    y = 1024  # 小块边长y
    overlap = int(1024*0.1)  # 重叠区域
    split_images_and_masks(data_root, output_folder, x, y, overlap)